import sys
from copy import deepcopy
from pyevolve import GSimpleGA
from pyevolve import Consts
from pyevolve import GTree
import game
import antPickel

## Constancts
# The maximum depth of a genome tree
MAX_TREE_DEPTH = 3

# The genome method
GENOME_MTH = "full"

# This number of times will next_generation() of generator.py be called:
MAX_GENERATION = 100

# The number of steps each ant can make before we declare the generation 
# successful/unsuccessful:
MAX_STEPS = game.CELL_CNT * 3

# The crossover rate to use in the GA
CROSSOVER_RATE = .9

# The probablity that a mutation happens
MUTATION_RATE = 0.1

# The number of ant colonies where a colonie is a groups of ants, where each 
# ant in the group has the same chromosome but groups can (and hopfuly do)
# have diffrent chromosomes
NUM_ANT_COLONIES = 150
    
## Ant Gentric Programing Controler
class AntGP:
    ''' A class to control the Gentic Algorithm.
    '''
    def __init__(   self,
                    max_gen = MAX_GENERATION, co_rate = CROSSOVER_RATE,
                    mut_rate = MUTATION_RATE, pop = NUM_ANT_COLONIES, 
                    max_tree_depth = MAX_TREE_DEPTH, mtd = GENOME_MTH):
        ''' Initalze the class with the given GA constants.
        '''
        genome = GTree.GTreeGP()
        genome.setParams(max_depth=max_tree_depth, method=mtd)
        
        # Set eval functions
        def s_w_func(f, scale, weight): 
            return lambda ant: scale_and_weight(f(ant), scale, weight)
        #genome.evaluator += s_w_func(evalfun_morsels_left,game.MORSEL_CNT,7)
        #genome.evaluator += s_w_func(evalfun_path_coverage,game.CELL_CNT,2)
        #genome.evaluator += s_w_func(evalfun_spining,MAX_STEPS,1)
        genome.evaluator += evalfun_fast
        
        # Create ga
        self.ga = GSimpleGA.GSimpleGA(genome)
        self.ga.setParams(gp_terminals  = ['advanced', 'turnright', 'turnleft'],
                    gp_function_prefix  = "antfun")
    
        self.ga.setMinimax(Consts.minimaxType["minimize"])
        self.ga.setGenerations(max_gen)
        self.ga.setCrossoverRate(co_rate)
        self.ga.setMutationRate(mut_rate)
        self.ga.setPopulationSize(pop)

        self.has_ran = False
        
    def run_ga(self, stats_at = 0):
        ''' Run the ga and print stats at the given interaval, in generations.
        '''
        if stats_at == 0:
            self.ga()
        else:
            self.ga(freq_stats=stats_at)
        self.has_ran = True
        
    def best_individual(self):
        ''' Runs the ga (if not alerady done) and returns the best individual
        found.
        '''
        if(not self.has_ran): run_ga()
        return self.ga.bestIndividual()
    
## GP Functions
# GP Inner Node Functions
def antfun_if_sensor_forward(ifTrue, ifFalse):
    ''' GA Function, checks if there is a morsel ahead. Returns a function
    that takes a ant to perform the check on. If true, perfroms the first
    arguement else perfroms the second.
    '''
    trueVal = lambda ant: ifTrue(ant)
    falseVal = lambda ant: ifFalse(ant)
    return lambda ant: trueVal(ant) if ant.is_morsel_ahead() else falseVal(ant)

def antfun_if_sensor_right(ifTrue, ifFalse):
    ''' GA Function, checks if there is a morsel to the right. Returns a 
    function that takes a ant to perform the check on. If true, perfroms the 
    first arguement else perfroms the second.
    '''
    trueVal = lambda ant: ifTrue(ant)
    falseVal = lambda ant: ifFalse(ant)
    return lambda ant: trueVal(ant) if ant.is_morsel_right() else falseVal(ant)
 
def antfun_if_sensor_left(ifTrue, ifFalse):
    ''' GA Function, checks if there is a morsel to the left. Returns a 
    function that takes a ant to perform the check on. If true, perfroms the 
    first arguement else perfroms the second.
    '''
    trueVal = lambda ant: ifTrue(ant)
    falseVal = lambda ant: ifFalse(ant)
    return lambda ant: trueVal(ant) if ant.is_morsel_left() else falseVal(ant)
     
def antfun_progn2(first, second):
    ''' GA Function, used to combine two other GAFunctions. Returns a function
    that takes a ant to perform on the two inputs.
    '''
    def do_both(ant):
        first(ant)
        second(ant)
    return do_both
   
def antfun_progn3(first, second, third):
    ''' GA Function, used to combine three other GAFunctions. Returns a function
    that takes a ant to perform on the three inputs.
    '''
    def do_all(ant):
        first(ant)
        second(ant)
        third(ant)
    return do_all
   
def antfun_progn4(first, second, third, fouth):
    ''' GA Function, used to combine four other GAFunctions. Returns a function
    that takes a ant to perform on the four inputs.
    '''
    def do_all(ant):
        first(ant)
        second(ant)
        third(ant)
        fouth(ant)
    return do_all
    
# GP Leaf Node Functions
antterm_advanced = lambda ant: ant.advance()
antterm_turnright = lambda ant: ant.turn_right()
antterm_turnleft = lambda ant: ant.turn_left()

## Eval functions
def evalfun_morsels_left(chromosome):
    ''' Run ant game for the given chromosome. Return the number of morsels.
    '''    
    # Run GP chromosome
    board = peform_chromosome(chromosome)
    
    # The fitness score is the number of morsels that are left on the gameboard
    return board.num_morsels
    
def evalfun_path_coverage(chromosome):
    ''' Run ant game for the given chromosome. Return the number of positions on
    the board not visited.
    '''
    # Make counters
    pos_counts = {}
    def check_ant_locations(board):
        for ant in board.ants_in():
            pos_counts[ant.ID] = True
            #if ant.ID in pos_counts:
            #    pos_counts[ant.ID] += 1
            #else:
            #    pos_counts[ant.ID] = 1
    
    # Run GP chromosome
    board = peform_chromosome(chromosome, check_ant_locations)
    
    # The fitness score is the number of positions not entered
    return game.CELL_CNT - len(pos_counts)
 
def evalfun_spining(chromosome):
    ''' Run ant game for the given chromosome. Returns the number times the ant 
    stays in the same place.
    '''
    counts = {}
    counts["lastID"] = None
    counts["lastCount"] = None
    counts["totalCount"] = 0
    def check_ant_spinning(board):
        for ant in board.ants_in():
            add_one = False
            if counts["lastID"] == ant.ID:
                counts["lastCount"] += 1
                if counts["lastCount"] > 4:
                    add_one = True
            else:
                counts["lastID"] = ant.ID
                counts["lastCount"] = 1
            if add_one:
                counts["totalCount"] += 1
                    
    # Run GP chromosome
    board = peform_chromosome(chromosome, check_ant_spinning)
    
    # The fitness score is the number times the ant stays in the same place
    return counts["totalCount"]
 
def evalfun_fast(chromosome):
    '''
    '''
    # Make counters
    pos_counts = {}
    counts = {}
    counts["lastID"] = None
    counts["lastCount"] = None
    counts["totalCount"] = 0
    def fast_check(board):
        for ant in board.ants_in():
            pos_counts[ant.ID] = True
            
            add_one = False
            if counts["lastID"] == ant.ID:
                counts["lastCount"] += 1
                if counts["lastCount"] > 4:
                    add_one = True
            else:
                counts["lastID"] = ant.ID
                counts["lastCount"] = 1
            if add_one:
                counts["totalCount"] += 1
                
    # Run GP chromosome
    board = peform_chromosome(chromosome, fast_check)
    
    # The fitness score is the number of morsels that are left on the gameboard
    num_morsels = board.num_morsels
    
    # The fitness score is the number of positions not entered
    coverage =  game.CELL_CNT - len(pos_counts)
    
    # The fitness score is the number times the ant stays in the same place
    spining = counts["totalCount"]
    
    num_morsels = scale_and_weight(num_morsels,game.MORSEL_CNT,7)
    coverage = scale_and_weight(coverage, game.CELL_CNT,2)
    spining = scale_and_weight(spining, MAX_STEPS,1)
    
    return num_morsels + coverage + spining
    
## Helper Functions
def scale_and_weight(val, scale, weight):
    ''' Scale and weigh the given value. Reutrned value is a double.
    '''
    return weight * (1.00 * val / scale)
       
def num_moves(chromosome):
    ''' Returns the number of leaf nodes (move forward, turn left, turn right).
    '''
    moves = 0
    for node in chromosome.getAllNodes():
        if chromosome.getNodeHeight(node) == 0: # Check if leaf node, or terimal
            moves += 1
    return moves
    
def make_antfun(chromosome, callback_func = None):
    ''' Return a function that takes an ant and when called perfroms the given 
    chromosome. If there is more then 1 argument, the given function is called
    at each terminal. The function is given the current board as an argument.
    '''
    # Create terminals
    advanced = None
    turnright = None
    turnleft = None
    
    if callback_func is None:
        advanced = antterm_advanced
        turnright = antterm_turnright
        turnleft = antterm_turnleft
    else:
        def make_both_func(org_fun):
            def both_func(ant):
                org_fun(ant)
                callback_func(ant.board)
            return both_func
            
        advanced = make_both_func(antterm_advanced) 
        turnright = make_both_func(antterm_turnright) 
        turnleft = make_both_func(antterm_turnleft) 
    
    # Run GP chromosome
    code_comp = chromosome.getCompiledCode()
    return eval(code_comp)
    
def peform_chromosome(chromosome, callback_func = None):
    ''' Perfroms the chromosme on a new board. Returns the board after the
    chromosome was performed. The callback_fun is called after each move given
    the value of the board.
    '''
    board = game.get_board()
    chrom_func = make_antfun(chromosome, callback_func)
    moves_left = MAX_STEPS
    moves_per_chromosome = num_moves(chromosome)
    while moves_left > 0:
	   complete = board.peform_chromosome(chrom_func)
	   if complete: return board
		
	   moves_left -= moves_per_chromosome
    return board

def get_boards_for_chromosome(chromosome):
    ''' Return a list of the board from the given chromosome.
    '''
    boards = []
    get_boards = lambda board: boards.append(deepcopy(board))
    peform_chromosome(chromosome, get_boards)
    
    return boards

def main_run():
    
    # Run GA
    antProblem = AntGP()
    antProblem.run_ga(1)
    
    # Get best chromosome
    bestChrom = antProblem.best_individual()
    print "Best:"
    print bestChrom
    
    # Save chromosome
    file_name = None
    if len(sys.argv) == 1:
        file_name = "best.chrm"
    else:
        file_name =  sys.argv[1]
    antPickel.pickle_chromosome(bestChrom,game.CURR_SEED,file_name)
    print "(saved as", file_name, ")"
    
    # Print Boards
    print "Start:"
    board = game.get_board()
    print board
    print "End"
    board = peform_chromosome(bestChrom)
    print board
    
    # Print Stats
    print "Missed:", evalfun_morsels_left(bestChrom), "out of", game.MORSEL_CNT
    print "Positions Missed:", evalfun_path_coverage(bestChrom)
    print "Time Spining:", evalfun_spining(bestChrom)

if __name__ == "__main__":
    main_run()
